package com.zws.cucumber.optimization.socket.oio;

import com.zws.cucumber.optimization.socket.connector.ZwsConnectorAdaptor;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.util.ReflectionUtils;

import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Method;
import java.net.*;
import java.util.LinkedHashMap;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

import static com.zws.cucumber.optimization.socket.ZwsSockets.isMainThread;

public class ZwsSocketImplAdapter extends SocketImpl {

    private static final Logger logger = LoggerFactory.getLogger(ZwsSocketImplAdapter.class);

    private static final Map<Integer, ZwsConnectorAdaptor> port2Connectors = new ConcurrentHashMap<>();

    private static final Method m_create_boolean;
    private static final Method m_connect_stringAndInt;
    private static final Method m_connect_inetAddressAndInt;
    private static final Method m_connect_socketAddressAndInt;
    private static final Method m_getInputStream;
    private static final Method m_getOutputStream;
    private static final Method m_available;
    private static final Method m_close;
    private static final Method m_sendUrgentData;

    static {
        m_create_boolean = ReflectionUtils.findMethod(SocketImpl.class, "create", boolean.class);
        m_create_boolean.setAccessible(true);
        m_connect_stringAndInt = ReflectionUtils.findMethod(SocketImpl.class, "connect", String.class, int.class);
        m_connect_stringAndInt.setAccessible(true);
        m_connect_inetAddressAndInt = ReflectionUtils.findMethod(SocketImpl.class, "connect", InetAddress.class, int.class);
        m_connect_inetAddressAndInt.setAccessible(true);
        m_connect_socketAddressAndInt = ReflectionUtils.findMethod(SocketImpl.class, "connect", SocketAddress.class, int.class);
        m_connect_socketAddressAndInt.setAccessible(true);

        m_getInputStream = ReflectionUtils.findMethod(SocketImpl.class, "getInputStream");
        m_getInputStream.setAccessible(true);
        m_getOutputStream = ReflectionUtils.findMethod(SocketImpl.class, "getOutputStream");
        m_getOutputStream.setAccessible(true);

        m_available = ReflectionUtils.findMethod(SocketImpl.class, "available");
        m_available.setAccessible(true);

        m_close = ReflectionUtils.findMethod(SocketImpl.class, "close");
        m_close.setAccessible(true);

        m_sendUrgentData = ReflectionUtils.findMethod(SocketImpl.class, "sendUrgentData", int.class);
        m_sendUrgentData.setAccessible(true);
    }

    boolean isOptimized = false;
    boolean isDefaultConnected = false;
    volatile boolean isConnected = false;
    SocketImpl impl;

    SocketImpl defaultImpl;
    SocketImpl memoryImpl;

    ConnectFunc delayedConnectFunc;

    InputStream defaultInputStream;
    OutputStream defaultOutputStream;

    InputStream adapterInputStream;
    OutputStream adapaterOutputStream;

    boolean stream;
    Map<Integer, Object> options = new LinkedHashMap<>();

    public ZwsSocketImplAdapter() {
        adapterInputStream = new ZwsSocketImplAdapterInputStream(this);
        adapaterOutputStream = new ZwsSocketImplAdapterOutputStream(this);
    }

    @Override
    protected void create(boolean stream) throws IOException {
//        logger.debug("zws, create");
        // noop
        this.stream = stream;
    }

    @Override
    protected void connect(String host, int port) throws IOException {
//        logger.debug("zws, connect");
        logger.debug("connect(String, int), {}, {}", host, port);
        if (("127.0.0.1".equals(host) || "localhost".equals(host)) && port2Connectors.containsKey(port)) {
            isOptimized = true;

            initMemorySocketImpl(port2Connectors.get(port));
            if (isMainThread()) {
                delayedConnectFunc = () -> invoke(m_connect_stringAndInt, defaultImpl(), host, port);
            } else {
                connect0(() -> invoke(m_connect_stringAndInt, defaultImpl(), host, port));
            }

        } else {
            connect0(() -> invoke(m_connect_stringAndInt, defaultImpl(), host, port));
        }
        isConnected = true;
    }

    @Override
    protected void connect(InetAddress address, int port) throws IOException {
//        logger.debug("zws, connect");
        logger.debug("connect(InetAddress, int), {}, {}, {}, {}", address, port, address.isAnyLocalAddress(), address.isLoopbackAddress());
        if (address.isLoopbackAddress() && port2Connectors.containsKey(port)) {
            isOptimized = true;

            initMemorySocketImpl(port2Connectors.get(port));
            if (isMainThread()) {
                delayedConnectFunc = () -> invoke(m_connect_inetAddressAndInt, defaultImpl(), address, port);
            } else {
                connect0(() -> invoke(m_connect_inetAddressAndInt, defaultImpl(), address, port));
            }

        } else {
            connect0(() -> invoke(m_connect_inetAddressAndInt, defaultImpl(), address, port));
        }
        isConnected = true;
    }

    @Override
    protected void connect(SocketAddress addr, int timeout) throws IOException {
//        logger.debug("zws, connect");
        logger.debug("connect(SocketAddress, int), {}, {}", addr, addr instanceof InetSocketAddress);
        if (addr instanceof InetSocketAddress addr0) {
            logger.debug("connect(SocketAddress, int), {}, {}, {}", addr0.getPort(), addr0.getAddress().isAnyLocalAddress(), addr0.getAddress().isLoopbackAddress());
        }
        if (addr instanceof InetSocketAddress addr0
                && addr0.getAddress().isLoopbackAddress()
                && port2Connectors.containsKey(addr0.getPort())) {
            isOptimized = true;

            initMemorySocketImpl(port2Connectors.get(addr0.getPort()));
            if (isMainThread()) {
                delayedConnectFunc = () -> invoke(m_connect_socketAddressAndInt, defaultImpl(), addr, timeout);
            } else {
                connect0(() -> invoke(m_connect_socketAddressAndInt, defaultImpl(), addr, timeout));
            }

        } else {
            connect0(() -> invoke(m_connect_socketAddressAndInt, defaultImpl(), addr, timeout));
        }
        isConnected = true;
    }

    private void connect0(ConnectFunc func) throws IOException {
        logger.debug("use default socket impl, stream: {}", stream);
        initDefaultSocketImpl();

        invoke(m_create_boolean, defaultImpl(), stream);
        for (Map.Entry<Integer, Object> en : options.entrySet()) {
//            logger.debug("zws: set option, key: {}, val: {}", getOptionKey(en.getKey()), en.getValue());
            defaultImpl().setOption(en.getKey(), en.getValue());
        }
        func.apply();

        defaultInputStream = (InputStream) invoke(m_getInputStream, defaultImpl());
        defaultOutputStream = (OutputStream) invoke(m_getOutputStream, defaultImpl());
//        logger.debug("zws: use default socket impl, finished");
    }

    @Override
    protected void bind(InetAddress host, int port) throws IOException {
//        logger.debug("zws, bind");
        throw new UnsupportedOperationException();
    }

    @Override
    protected void listen(int backlog) throws IOException {
//        logger.debug("zws, listen");
        throw new UnsupportedOperationException();
    }

    @Override
    protected void accept(SocketImpl s) throws IOException {
//        logger.debug("zws, accept");
        throw new UnsupportedOperationException();
    }

    @Override
    protected InputStream getInputStream() throws IOException {
//        logger.debug("zws, getInputStream");
//        if (isDefaultConnected()) {
//            return (InputStream) invoke(m_getInputStream, impl);
//        } else {
//            return memoryImpl().getInputStream();
//        }
        return adapterInputStream;
    }

    @Override
    protected OutputStream getOutputStream() throws IOException {
//        logger.debug("zws, getOutputStream");
//        if (isDefaultConnected()) {
//            return (OutputStream) invoke(m_getOutputStream, impl());
//        } else {
//            return memoryImpl().getOutputStream();
//        }
        return adapaterOutputStream;
    }

    @Override
    protected int available() throws IOException {
//        logger.debug("zws: available");
        if (isOptimized() && isMainThread()) {
            return memoryImpl().available();
        }
        return (int) invoke(m_available, defaultImpl());
    }

    @Override
    protected void close() throws IOException {
//        logger.debug("zws: close, isDefault: {}", isDefault());
        if (isOptimized()) {
            memoryImpl().close();
        }
        if (isDefaultConnected()) {
            invoke(m_close, defaultImpl());
        }
    }

    @Override
    protected void sendUrgentData(int data) throws IOException {
//        logger.debug("zws, sendUrgentData");
        if (isOptimized() && isMainThread()) {
            memoryImpl().sendUrgentData(data);
        }
        invoke(m_sendUrgentData, defaultImpl(), data);
    }

    @Override
    public void setOption(int optID, Object value) throws SocketException {
//        logger.debug("zws: set option, key: {}, val: {}", getOptionKey(optID), value);
        if (!isConnected) {
            options.put(optID, value);

        } else {
            if (isOptimized() && isMainThread()) {
                options.put(optID, value);
            } else {
                defaultImpl().setOption(optID, value);
            }
        }
    }

    @Override
    public Object getOption(int optID) throws SocketException {
//        logger.debug("zws: get option, key: {}", getOptionKey(optID));
        if (!isConnected) {
            return options.get(optID);
        } else {
            if (isOptimized() && isMainThread()) {
                return options.get(optID);
            } else {
                return defaultImpl().getOption(optID);
            }
        }
    }

    protected void initMemorySocketImpl(ZwsConnectorAdaptor connectorAdaptor) {
        logger.debug("use memory socket impl");
        memoryImpl(new ZwsSocketImpl(connectorAdaptor));
    }

    protected void initDefaultSocketImpl() {
        defaultImpl(createDefaultSocketImpl());
        this.isDefaultConnected = true;
    }

    private void impl(SocketImpl impl) {
        assert this.impl == null;
        this.impl = impl;
    }

    private SocketImpl impl() {
        return impl;
    }

    protected SocketImpl activeImpl() {
        if (isOptimized()) {
            return memoryImpl();
        }
        return defaultImpl();
    }

    protected InputStream activeInputStream() throws IOException {
        if (isOptimized()) {
            return memoryImpl().getInputStream();
        }
        return defaultInputStream;
    }

    protected OutputStream activeOutputStream() throws IOException {
        if (isOptimized()) {
            return memoryImpl().getOutputStream();
        }
        return defaultOutputStream;
    }

    protected void defaultImpl(SocketImpl impl) {
        this.defaultImpl = impl;
    }

    protected SocketImpl defaultImpl() {
        /*
         * if default is null, that means memoryImpl is used(connected to jetty or wiremock)
         * but this socket is reused in other thread
         */
        if (defaultImpl == null) {
            if (isOptimized()) {
                try {
                    connect0(delayedConnectFunc);
                } catch (IOException e) {
                    throw new RuntimeException(e);
                }
            } else {
                throw new IllegalStateException("defaultImpl should not be null here if it is not optimized.");
            }
        }
        return defaultImpl;
    }

    protected void memoryImpl(ZwsSocketImpl impl) {
        this.memoryImpl = impl;
    }

    protected ZwsSocketImpl memoryImpl() {
        return (ZwsSocketImpl) memoryImpl;
    }

    public boolean isOptimized() {
        return isOptimized;
    }

    private static SocketImpl createDefaultSocketImpl() {
        try {
            Class<?> clazz = Class.forName("java.net.SocketImpl");
            Method method = clazz.getDeclaredMethod("createPlatformSocketImpl", boolean.class);
            method.setAccessible(true);
            return (SocketImpl) method.invoke(clazz, false);

        } catch (ClassNotFoundException
                 | NoSuchMethodException
                 | IllegalAccessException
                 | IllegalArgumentException
                 | InvocationTargetException e) {
            throw new RuntimeException(e);
        }
    }

    private static Object invoke(Method method, Object impl, Object... args) throws IOException {
        try {
            return method.invoke(impl, args);
        } catch (InvocationTargetException e) {
            Throwable cause = e.getCause();
            if (cause instanceof IOException ioex) {
                throw ioex;
            } else {
                throw (RuntimeException) cause;
            }
        } catch (IllegalAccessException | IllegalArgumentException e) {
            throw new RuntimeException(e);
        }
    }

    private boolean isDefaultConnected() {
        return this.isDefaultConnected;
    }

    public static void registerLocalConnectorAdaptor(int port, ZwsConnectorAdaptor adaptor) {
        port2Connectors.put(port, adaptor);
    }

    @FunctionalInterface
    public interface ConnectFunc {
        void apply() throws IOException;
    }

    private String getOptionKey(int opt) {
        return switch (opt) {
            case 0x0001 -> "TCP_NODELAY";
            case 0x000F -> "SO_BINDADDR";
            case 0x04 -> "SO_REUSEADDR";
            case 0x0E -> "SO_REUSEPORT";
            case 0x0020 -> "SO_BROADCAST";
            case 0x10 -> "IP_MULTICAST_IF";
            case 0x1f -> "IP_MULTICAST_IF2";
            case 0x12 -> "IP_MULTICAST_LOOP";
            case 0x3 -> "IP_TOS";
            case 0x0080 -> "SO_LINGER";
            case 0x1006 -> "SO_TIMEOUT";
            case 0x1001 -> "SO_SNDBUF";
            case 0x1002 -> "SO_RCVBUF";
            case 0x0008 -> "SO_KEEPALIVE";
            case 0x1003 -> "SO_OOBINLINE";
            default -> throw new IllegalArgumentException(Integer.toHexString(opt));
        };
    }
}
